/*

BranchMonitor - ROP Mitigation using block entry point identification as a heuristic
Richard Johnson <rjohnson@sourcefire.com>

*/
#include "pin.h"
#include <stdio.h>
#include <process.h>
#include <stdint.h>
#include <ctime>
#include <stack>
#include <set>

set<ADDRINT> ropBlocks;

clock_t startclock;

#pragma region System handled events

BOOL ChildProcessCallback(CHILD_PROCESS childProcess, VOID *val)
{
	LOG("FORK: Parent Process: " + decstr(_getpid()) + " Child Process: " + decstr(CHILD_PROCESS_GetId(childProcess)) + "\n");

	// inject into child process
	return TRUE;
}  

VOID Fini(INT32 code, VOID *v) 
{

}

void ModuleLoadCallback(IMG image, VOID *v) 
{
	ADDRINT imgLowAddr = IMG_LowAddress(image);
	ADDRINT imgHighAddr = IMG_HighAddress(image);
	
	LOG("[+] Module image loaded: [" + 
		hexstr(imgLowAddr, sizeof(imgLowAddr) * 2) + 
		" - " + 
		hexstr(imgHighAddr, sizeof(imgHighAddr) * 2) + "] " + 
		IMG_Name(image) + "\n");

	return;
}

VOID ModuleUnloadCallback(IMG image, VOID *v) 
{
	LOG("[+] Module image unloaded: " + IMG_Name(image) + "\n");
}

#pragma endregion

#pragma region Threads

void ThreadStart(THREADID tid, CONTEXT * context, int flags, void * v)
{ 
	LOG("[+] Thread Create [tid=" + decstr(tid) + "]\n");
}

void ThreadFini(THREADID tid, const CONTEXT * context, int code, void * v)
{
	LOG("[+] Thread Exit [tid=" + decstr(tid) + "] code: " + hexstr(code, sizeof(code) * 2) + "\n");
}
#pragma endregion


#pragma region Exception Handling
VOID ExceptionCallback(THREADID tid, CONTEXT_CHANGE_REASON reason, const CONTEXT *from, CONTEXT *context, INT32 info, VOID *v)
{
	if (reason != CONTEXT_CHANGE_REASON_EXCEPTION)
		return;

	uint32_t eip = PIN_GetContextReg(from, REG_EIP);
	uint32_t esp = PIN_GetContextReg(from, REG_ESP);
	uint32_t ebp = PIN_GetContextReg(from, REG_EBP);
	uint32_t childebp = 0;
	
	char buf[1024];
	memset(buf, 0, sizeof(buf));
	_snprintf(buf, sizeof(buf) - 1, 
		"\nException code=0x%x address=0x%x tid=%d\n\n"
		"eax=%08x ebx=%08x ecx=%08x edx=%08x esi=%08x edi=%08x\n"
		"eip=%08x esp=%08x ebp=%08x\n",
		info, eip, tid,
		PIN_GetContextReg(from, REG_EAX),
		PIN_GetContextReg(from, REG_EBX),
		PIN_GetContextReg(from, REG_ECX),
		PIN_GetContextReg(from, REG_EDX),
		PIN_GetContextReg(from, REG_ESI),
		PIN_GetContextReg(from, REG_EDI),
		eip, esp, ebp);		
	LOG(buf);

	// log callstack		
	memset(buf, 0, sizeof(buf));
	_snprintf(buf, sizeof(buf) - 1, 
		"\nCallstack:\n"
		" FramePtr ChildEBP RetAddr\n");
	LOG(buf);

	int count = 0;
	memset(buf, 0, sizeof(buf));
	while(ebp != 0 && count < 20)
	{
		if(PIN_SafeCopy(&childebp, (uint32_t *)(ebp), 4) != 4) 
			break;
		if(PIN_SafeCopy(&eip, (uint32_t *)(ebp + 4), 4) != 4) 
			break;		

		_snprintf(buf, sizeof(buf) -1, "%08x %08x %08x\n", 
			ebp, childebp, eip);
		LOG(buf);

		if(PIN_SafeCopy(&ebp, (uint32_t *)ebp, 4) != 4) 
			break;

		count++;
	}

	time_t date = time(0);	
	LOG("\n Trace finished in " + decstr((clock() - startclock) / CLOCKS_PER_SEC) + " seconds\n");
	LOG(ctime(&date));

	if(info == 0xC0000005)
	{				
		exit(1);
	}
}
#pragma endregion


#pragma region Control flow events

void RetHook(THREADID tid, ADDRINT address, CONTEXT *context)
{
	// Mitigations
	uint32_t savedEip = 0;
	uint32_t shadowEip = 0;
	uint32_t stackPtr = (uint32_t)PIN_GetContextReg(context, REG_ESP);
	PIN_SafeCopy(&savedEip, (uint32_t *)stackPtr, 4);		

	//// ROP Block	
	if(ropBlocks.find(savedEip) == ropBlocks.end())
	{
		LOG("\n*** ROP BLOCK ERROR: savedEip[" + hexstr(savedEip, 4) + "] is not a known block entry point ***\n");
	}

}
#pragma endregion

#pragma region Instrumentation

VOID RoutineCallback(RTN routine, VOID *v) 
{

	ADDRINT addr = RTN_Address(routine);
	RTN_Open(routine);

	//Mitigations
	ropBlocks.insert(addr);
	for (BBL bbl = RTN_BblHead(routine); BBL_Valid(bbl); bbl = BBL_Next(bbl)) 
	{
		ropBlocks.insert(BBL_Address(bbl));
		for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins=INS_Next(ins))
		{
				if(INS_IsCall(ins))
					ropBlocks.insert(INS_Address(ins) + INS_Size(ins));
		}

	}
	RTN_Close(routine);
}

VOID TraceCallback(TRACE trace, VOID *v) 
{
	ADDRINT addr = TRACE_Address(trace);
	char buf[256];
	memset(buf, 0, 256);

	ropBlocks.insert(addr);
	for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl)) 
	{
		ropBlocks.insert(BBL_Address(bbl));
		for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins=INS_Next(ins))
		{
			if(INS_IsCall(ins))
			{
				int insSize = INS_Size(ins);
				ropBlocks.insert(INS_Address(ins) + insSize);

			}
			else if(INS_IsRet(ins))
			{
				INS_InsertCall(ins, IPOINT_BEFORE, AFUNPTR(RetHook),																		
										IARG_THREAD_ID,
										IARG_INST_PTR,
										IARG_CONTEXT,
										IARG_END);					
			}
		}
	}
}
#pragma endregion

int main(int argc, char *argv[]) 
{
	// Initialize the PIN engine
	PIN_InitSymbols();
	
	if(PIN_Init(argc, argv))
	{
		return 1;
	}
	time_t date = time(0);	
	LOG("[+] " + string(ctime(&date)));
	startclock = clock();

	PIN_AddFollowChildProcessFunction(ChildProcessCallback, 0);
	
	PIN_AddFiniFunction(Fini, NULL);
	PIN_AddThreadStartFunction(ThreadStart, NULL);
	PIN_AddThreadFiniFunction(ThreadFini, NULL);
	
	IMG_AddInstrumentFunction(ModuleLoadCallback, NULL);
	IMG_AddUnloadFunction(ModuleUnloadCallback, NULL);

	TRACE_AddInstrumentFunction(TraceCallback, NULL);
	RTN_AddInstrumentFunction(RoutineCallback, NULL);

	LOG("[+] Instrumentation hooks installed\n");

	PIN_AddContextChangeFunction(ExceptionCallback, 0);		

	date = time(0);	
	LOG("\n Trace beginning at " + decstr((clock() - startclock) / CLOCKS_PER_SEC) + " seconds\n");
	LOG(ctime(&date));

	PIN_StartProgram();

	return 0;
}
